{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torchvision import datasets, models, transforms\n",
    "import os\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_dir = \"DogsVSCats\"\n",
    "\n",
    "data_transform = {x:transforms.Compose([transforms.Scale([224,224]),\n",
    "                                        transforms.ToTensor(),\n",
    "                                        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5])])\n",
    "                  for x in [\"train\", \"valid\"]}\n",
    "\n",
    "image_datasets = {x:datasets.ImageFolder(root = os.path.join(data_dir,x),\n",
    "                                         transform = data_transform[x])\n",
    "                  for x in [\"train\", \"valid\"]}\n",
    "\n",
    "dataloader = {x:torch.utils.data.DataLoader(dataset= image_datasets[x],\n",
    "                                            batch_size = 16,\n",
    "                                            shuffle = True)\n",
    "              for x in [\"train\", \"valid\"]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "X_example, y_example = next(iter(dataloader[\"train\"]))\n",
    "example_clasees = image_datasets[\"train\"].classes\n",
    "index_classes = image_datasets[\"train\"].class_to_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model_1 = models.vgg16(pretrained=True)\n",
    "model_2 = models.resnet50(pretrained=True)\n",
    "\n",
    "Use_gpu = torch.cuda.is_available()\n",
    "\n",
    "for parma in model_1.parameters():\n",
    "    parma.requires_grad = False\n",
    "\n",
    "model_1.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 4096),\n",
    "                                             torch.nn.ReLU(),\n",
    "                                             torch.nn.Dropout(p=0.5),\n",
    "                                             torch.nn.Linear(4096, 4096),\n",
    "                                             torch.nn.ReLU(),\n",
    "                                             torch.nn.Dropout(p=0.5),\n",
    "                                             torch.nn.Linear(4096, 2))\n",
    "\n",
    "for parma in model_2.parameters():\n",
    "    parma.requires_grad = False\n",
    "\n",
    "model_2.fc = torch.nn.Linear(2048, 2)\n",
    "    \n",
    "if Use_gpu:\n",
    "    model_1 = model_1.cuda()\n",
    "    model_2 = model_2.cuda()\n",
    "    \n",
    "loss_f_1 = torch.nn.CrossEntropyLoss()\n",
    "loss_f_2 = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "optimizer_1 = torch.optim.Adam(model_1.classifier.parameters(), lr =0.00001)\n",
    "optimizer_2 = torch.optim.Adam(model_2.fc.parameters(), lr = 0.00001)\n",
    "\n",
    "weight_1 = 0.6\n",
    "weight_2 = 0.4\n",
    "epoch_n = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "time_open = time.time()\n",
    "for epoch in range(epoch_n):\n",
    "    print(\"Epoch {}/{}\".format(epoch, epoch_n - 1))\n",
    "    for phase in [\"train\", \"valid\"]:\n",
    "        if phase == \"train\":\n",
    "            print(\"Training...\")\n",
    "            model_1.train(True)\n",
    "            model_2.train(True)\n",
    "        else:\n",
    "            print(\"Validing...\")\n",
    "            model_1.train(False)\n",
    "            model_2.train(False)\n",
    "            \n",
    "        running_loss_1 = 0.0\n",
    "        running_corrects_1 = 0\n",
    "        running_loss_2 = 0.0\n",
    "        running_corrects_2 = 0\n",
    "        blending_running_corrects = 0\n",
    "    \n",
    "        for batch, data in enumerate(dataloader[phase], 1):\n",
    "            X, y = data\n",
    "            \n",
    "            if Use_gpu:\n",
    "                X, y = Variable(X.cuda()), Variable(y.cuda())\n",
    "            else:\n",
    "                X, y = Variable(X), Variable(y)\n",
    "                \n",
    "            y_pred_1 = model_1(X)\n",
    "            y_pred_2 = model_2(X)\n",
    "            blending_y_pred = y_pred_1*weight_1+y_pred_2*weight_2\n",
    "            \n",
    "            _, pred_1 = torch.max(y_pred_1.data, 1)\n",
    "            _, pred_2 = torch.max(y_pred_2.data, 1)\n",
    "            _, blending_pred = torch.max(blending_y_pred.data, 1)\n",
    "            optimizer_1.zero_grad()\n",
    "            optimizer_2.zero_grad()\n",
    "            \n",
    "            loss_1 = loss_f_1(y_pred_1, y)\n",
    "            loss_2 = loss_f_2(y_pred_2, y)\n",
    "            \n",
    "            if phase == \"train\":\n",
    "                loss_1.backward()\n",
    "                loss_2.backward()\n",
    "                optimizer_1.step()\n",
    "                optimizer_2.step()\n",
    "                \n",
    "            running_loss_1 += loss_1.data[0]\n",
    "            running_corrects_1 += torch.sum(pred_1 == y.data)\n",
    "            running_loss_2 += loss_2.data[0]\n",
    "            running_corrects_2 += torch.sum(pred_2 == y.data)\n",
    "            blending_running_corrects += torch.sum(blending_pred ==y.data)\n",
    "            \n",
    "            if batch%500 == 0 and phase ==\"train\":\n",
    "                print(\"Batch {},Model1 Train Loss:{:.4f},Model1 Train ACC:{:.4f},Model2 \\\n",
    "                Train Loss:{:.4f},Model2 Train ACC:{:.4f}, \\\n",
    "                Blending_Model ACC:{:.4f}\".format(batch,running_loss_1/batch,100*running_corrects_1/(16*batch),running_loss_2/batch,100*running_corrects_2/(16*batch),100*blending_running_corrects/(16*batch)))\n",
    "\n",
    "       \n",
    "        epoch_loss_1 = running_loss_1*16/len(image_datasets[phase])\n",
    "        epoch_acc_1 = 100*running_corrects_1/len(image_datasets[phase])\n",
    "        epoch_loss_2 = running_loss_2*16/len(image_datasets[phase])\n",
    "        epoch_acc_2 = 100*running_corrects_2/len(image_datasets[phase])\n",
    "        epoch_blending_acc = 100*blending_running_corrects/len(image_datasets[phase])\n",
    "        \n",
    "        print(\"Epoch, Model1 Loss:{:.4f}, Model1 Acc:{:.4f}%, Model2 Loss:{:.4f}, Model2 Acc:{:.4f}%,\\\n",
    "                Blending_Model ACC:{:.4f}\".format(epoch_loss_1,epoch_acc_1, epoch_loss_2,epoch_acc_2,epoch_blending_acc))\n",
    "\n",
    "\n",
    "time_end = time.ti\n",
    "print(time_end)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
